# -*- coding: utf-8 -*-

"""
绘制二叉树图形显示方法
"""

import math
import uuid

from PyQt4 import QtCore, QtGui

class Edge(QtGui.QGraphicsItem):
    """二叉树节点间连线"""
    Pi = math.pi
    TwoPi = 2.0 * Pi

    def __init__(self, sourceNode, destNode):
        
        super(Edge, self).__init__()

        # 箭头大小
        self.arrowSize = 10.0
        self.source = sourceNode
        self.dest = destNode
        self.sourcePoint = QtCore.QPointF(sourceNode.pos())
        self.destPoint = QtCore.QPointF(destNode.pos())
        
    def boundingRect(self):
        """override"""
        if not self.source or not self.dest:
            return QtCore.QRectF()

        penWidth = 1.0
        extra = (penWidth + self.arrowSize) / 2.0

        return QtCore.QRectF(self.sourcePoint,
                QtCore.QSizeF(self.destPoint.x() - self.sourcePoint.x(),
                        self.destPoint.y() - self.sourcePoint.y())
            ).normalized().adjusted(-extra, -extra, extra, extra)

    def paint(self, painter, option, widget):
        """override"""
        if not self.source or not self.dest:
            return

        # Draw the line itself.
        line = QtCore.QLineF(self.sourcePoint, self.destPoint)

        if line.length() == 0.0:
            return

        painter.setPen(QtGui.QPen(QtCore.Qt.black, 1, QtCore.Qt.SolidLine,
                QtCore.Qt.RoundCap, QtCore.Qt.RoundJoin))
        painter.drawLine(line)

        # Draw the arrows if there's enough room.
        angle = math.acos(line.dx() / line.length())
        if line.dy() >= 0:
            angle = Edge.TwoPi - angle

        destArrowP1 = self.destPoint + QtCore.QPointF(math.sin(angle - Edge.Pi / 3) * self.arrowSize,
            math.cos(angle - Edge.Pi / 3) * self.arrowSize)
        destArrowP2 = self.destPoint + QtCore.QPointF(math.sin(angle - Edge.Pi + Edge.Pi / 3) * self.arrowSize,
            math.cos(angle - Edge.Pi + Edge.Pi / 3) * self.arrowSize)

        painter.setBrush(QtCore.Qt.black)
        painter.drawPolygon(QtGui.QPolygonF([line.p2(), destArrowP1, destArrowP2]))

class Node(QtGui.QGraphicsItem):
    """二叉树节点"""
    def __init__(self, nodeValue):
        
        super(Node, self).__init__()

        self.newPos = QtCore.QPointF()
        self.nodeValue = nodeValue
        self.setCacheMode(QtGui.QGraphicsItem.DeviceCoordinateCache)
        self.setZValue(1)
        
    def boundingRect(self):
        """override"""
        adjust = 2.0
        return QtCore.QRectF(-10 - adjust, -10 - adjust, 33 + adjust,
                33 + adjust)

    def paint(self, painter, option, widget):
        """override"""
        gradient = QtGui.QRadialGradient(-3, -3, 10)
        if option.state & QtGui.QStyle.State_Sunken:
            gradient.setCenter(3, 3)
            gradient.setFocalPoint(3, 3)
            gradient.setColorAt(1, QtGui.QColor(QtCore.Qt.yellow).light(120))
            gradient.setColorAt(0, QtGui.QColor(QtCore.Qt.darkYellow).light(120))
        else:
            gradient.setColorAt(0, QtCore.Qt.yellow)
            gradient.setColorAt(1, QtCore.Qt.darkYellow)

        painter.setBrush(QtGui.QBrush(gradient))
        painter.setPen(QtGui.QPen(QtCore.Qt.black, 0))
        painter.drawEllipse(-10, -10, 20, 20)
        
        # 在节点上显示数值
        textRect = QtCore.QRectF(-10, -10, 30, 30)
        message = self.nodeValue
        font = painter.font()
        font.setBold(True)
        font.setPointSize(10)
        painter.setFont(font)
        painter.setPen(QtCore.Qt.black)
        painter.drawText(textRect, message)

class Dialog(QtGui.QGraphicsView):
    """画图面板"""   
    def __init__(self,  tree):
        """
        初始化面板，并绘制二叉树
        @param tree: 二叉树实体
        """
        super(Dialog, self).__init__()
        
        # 窗体宽、高
        self.width = 1000
        self.height = 400
        
        scene = QtGui.QGraphicsScene(self)
        scene.setItemIndexMethod(QtGui.QGraphicsScene.NoIndex)
        scene.setSceneRect(0, 0, self.width,  self.height)
        self.setScene(scene)
        self.setCacheMode(QtGui.QGraphicsView.CacheBackground)
        self.setViewportUpdateMode(QtGui.QGraphicsView.BoundingRectViewportUpdate)
        self.setRenderHint(QtGui.QPainter.Antialiasing)
        self.setTransformationAnchor(QtGui.QGraphicsView.AnchorUnderMouse)
        self.setResizeAnchor(QtGui.QGraphicsView.AnchorViewCenter)
        self.setMinimumSize(self.width,  self.height)
        self.setWindowTitle(QtGui.QApplication.translate("DrawTree", "Huffman Tree"))
        
        if tree != None and len(tree) is not 0:
            # 给每一个节点添加唯一ID，并获得树的深度
            tree[0] = [uuid.uuid1(), tree[0]]
            max = [0]
            self.traversal(tree, maxdep=max)
            dep = max[0]
            
            # 画出二叉树
            self.drawTree(scene, tree, stepX = 6 * pow(2, dep - 1), stepY = 30)

    def traversal(self, tree, curdep=0, maxdep=[0]):
        """前序遍历二叉树"""
        if len(tree) == 2:
            tree[0] = [uuid.uuid1(), tree[0]]
            if curdep > maxdep[0]:
                maxdep[0] = curdep
        else:
            childL, childR = tree[1], tree[2]
            if isinstance(childL[0], int):
                childL[0] = [uuid.uuid1(), childL[0]]
            if isinstance(childR[0], int):
                childR[0] = [uuid.uuid1(), childR[0]]
            self.traversal(childL, curdep + 1, maxdep)
            self.traversal(childR, curdep + 1, maxdep)

    def drawTree(self, scene, tree, depth = 1, \
                 parentPosX = 500,  parentPosY = 0, \
                 parentValue = "", parentid = "", \
                 table = {}, stepX = 200, stepY = 30):
        """
        绘制二叉树
        @param depth: 树的深度
        @param parentPosX: 绘制父节点的X轴位置
        @param parentPosY: 绘制父节点的Y轴位置
        @param parentValue: 父节点的权值
        @param parentid: 父节点GUID
        @param stepX: 相邻兄弟节点间距离
        @param stepY: 父子节点间距离
        """
        if len(tree) == 2:
            pass
        else:
            nodeParent = None
            # 判断是否是根节点
            if parentValue == "":
                nodeParent = Node(str(tree[0][1]))
                scene.addItem(nodeParent)
                table[tree[0][0]] = nodeParent
                nodeParent.setPos(parentPosX, depth * stepY)
            else:
                # 判断该父节点是否已经在table表中
                if parentid in table:
                    nodeParent = table[parentid]
                else:
                    nodeParent = Node(str(tree[0][1]))
                    scene.addItem(nodeParent)
                    table[tree[0][0]] = nodeParent
            
            childL = tree[1]
            nodeL = None
            # 判断左子节点是否是叶子节点
            if len(childL) == 2:
                nodeL = Node(str(childL[0][1][1]) + "\n" + str(childL[1]))
                scene.addItem(nodeL)
                nodeL.setPos(parentPosX - stepX, (depth + 1) * stepY)
                table[childL[0][0]] = nodeL
                scene.addItem(Edge(nodeParent, nodeL))
            else:
                # 判断左子节点是否已经在table表中
                if childL[0][0] in table:
                    nodeL = table[childL[0][0]]
                else:
                    nodeL = Node(str(childL[0][1]))
                    scene.addItem(nodeL)
                    nodeL.setPos(parentPosX - stepX, (depth + 1) * stepY)
                    table[childL[0][0]] = nodeL
                scene.addItem(Edge(nodeParent, nodeL))
            
            self.drawTree(scene, tree[1], depth + 1, nodeL.x(), nodeL.y(), childL[0][1], childL[0][0], table, stepX / 1.8, stepY)
            
            childR = tree[2]
            nodeR = None
            # 判断右子节点是否是叶子节点
            if len(childR) == 2:
                nodeR = Node(str(childR[0][1][1]) + "\n" + str(childR[1]))
                scene.addItem(nodeR)
                nodeR.setPos(parentPosX + stepX, (depth + 1) * stepY)
                table[childR[0][0]] = nodeR
                scene.addItem(Edge(nodeParent, nodeR))
            else:
                # 判断右子节点是否已经在table表中
                if childR[0][0] in table:
                    nodeR = table[childR[0][0]]
                else:
                    nodeR = Node(str(childR[0][1]))
                    scene.addItem(nodeR)
                    nodeR.setPos(parentPosX + stepX, (depth + 1) * stepY)
                    table[childR[0][0]] = nodeR
                scene.addItem(Edge(nodeParent, nodeR))
            
            self.drawTree(scene, tree[2], depth + 1, nodeR.x(), nodeR.y(), childR[0][1], childR[0][0], table, stepX / 1.8, stepY)
        
    def drawBackground(self, painter, rect):
        """重载绘制背景函数，在背景上写上提示信息。"""
        sceneRect = self.sceneRect()
        textRect = QtCore.QRectF(sceneRect.left() + 4, sceneRect.top() + 4,
                sceneRect.width() - 4, sceneRect.height() - 4)
        message = "zoom with the mouse wheel"
        font = painter.font()
        font.setBold(True)
        font.setPointSize(14)
        painter.setFont(font)
        painter.setPen(QtCore.Qt.black)
        painter.drawText(textRect, message)
        
    def _scaleView(self, scaleFactor):
        """缩放窗体"""
        factor = self.matrix().scale(scaleFactor, scaleFactor).mapRect(QtCore.QRectF(0, 0, 1, 1)).width()

        if factor < 0.07 or factor > 100:
            return

        self.scale(scaleFactor, scaleFactor)
    
    def wheelEvent(self, event):
        """重载鼠标齿轮滚动事件，进行窗体缩放"""
        self._scaleView(math.pow(2.0, -event.delta() / 240.0))


if __name__ == "__main__":
    import sys
    tree = [12, [7, [4, [2, [1, 87], [1, 72]], [2, [1, 33], [1, 32]]], [3, [2, 111], [1, 114]]], [5, [3, 108], [2, [1, 101], [1, 100]]]]

    app = QtGui.QApplication(sys.argv)
    dialog = Dialog(tree)
    dialog.show()
    sys.exit(app.exec_())

